#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-08-06


"""
【节点名称】：
    YOLOv11DetNode_Cuda
【依赖项安装】：
    pip install spirems
    pip install ultralytics
【订阅类型】：
    sensor_msgs::CompressedImage （输入图像）
【发布类型】：
    spirecv_msgs::2DTargets （检测结果）
    sensor_msgs::CompressedImage （可视化结果，需借助可视化工具）
    std_msgs::Boolean （如果输入节点是数据集，则检测完成发布该话题让输入节点继续工作）
【构造参数说明】：
    parameter_file (str): 全局参数文件
    sms_shutdown (bool): 是否接收全局关闭信号，如果需要长期后台执行，建议设置为False
    specified_input_topic (str): 指定输入的话题地址
    specified_output_topic (str): 指定输出的话题地址
    realtime_det (bool): 是否是实时检测器，设置为True可以降低延迟，但可能会产生丢帧
【节点参数】：
    confidence (float): 目标得分阈值
    nms_thresh (float): NMS后处理参数
    dataset_name (str): 数据集名称
    pt_model (str): 加载模型名称
【备注】：
    无
"""

from ultralytics import YOLO
import threading
from queue import Queue
import numpy as np
import time
import os
from typing import Union
from spirems import Subscriber, Publisher, cvimg2sms, sms2cvimg, def_msg, QoS, BaseNode, get_extra_args
from spirems.mod_helper import download_model
import argparse
import platform
from spirecv.algorithm.utils import calc_fov, calc_los_pos


def trans_det_results(det_results, camera_matrix, calib_wh, img_wh, objs_in_meter, roi=None):
    sms_results = def_msg('spirecv_msgs::2DTargets')

    sms_results["file_name"] = ""
    sms_results["height"] = img_wh[1]  # int(det_results.orig_shape[0])
    sms_results["width"] = img_wh[0]  # int(det_results.orig_shape[1])
    has_calib = False
    if calib_wh[0] > 0 and calib_wh[1] > 0:
        sms_results["fov_x"], sms_results["fov_y"] = calc_fov(camera_matrix, calib_wh)
        has_calib = True

    if roi is not None:
        sms_results["rois"] = [[roi[0], roi[1], roi[2] - roi[0], roi[3] - roi[1]]]

    sms_results["targets"] = []
    cls = det_results.boxes.cls.cpu().numpy().astype(np.float64)
    conf = det_results.boxes.conf.cpu().numpy().astype(np.float64)
    xywh = det_results.boxes.xywh.cpu().numpy().astype(np.float64)
    for i in range(len(cls)):
        ann = dict()
        name = det_results.names[int(cls[i])].strip()
        ann["category_name"] = name.replace(' ', '_').lower()
        ann["category_id"] = int(cls[i])
        ann["score"] = round(conf[i], 3)
        ann["bbox"] = [round(j, 3) for j in xywh[i].tolist()]
        ann["bbox"][0] = ann["bbox"][0] - ann["bbox"][2] / 2
        ann["bbox"][1] = ann["bbox"][1] - ann["bbox"][3] / 2
        if roi is not None:
            ann["bbox"][0] += roi[0]
            ann["bbox"][1] += roi[1]

        ann["cxy"] = [
            (ann["bbox"][0] + ann["bbox"][2] / 2.) / sms_results["width"], 
            (ann["bbox"][1] + ann["bbox"][3] / 2.) / sms_results["height"]
        ]
        if has_calib and name in objs_in_meter:
            ann["los"], ann["pos"] = calc_los_pos(
                camera_matrix, calib_wh, 
                ann["cxy"], [ann["bbox"][2], ann["bbox"][3]], 
                objs_in_meter[name]
            )
        sms_results["targets"].append(ann)

    return sms_results


class YOLOv11DetNode_Cuda(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        threading.Thread.__init__(self)
        sms_shutdown = True if sms_shutdown in ['True', 'true', '1', True] else False
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )
        self.launch_next_emit = self.get_param("launch_next_emit", True)
        self.specified_input_topic = self.get_param("specified_input_topic", "")
        self.specified_output_topic = self.get_param("specified_output_topic", "")
        self.realtime_det = self.get_param("realtime_det", False)
        self.remote_ip = self.get_param("remote_ip", "127.0.0.1")
        self.remote_port = self.get_param("remote_port", 9094)
        self.confidence = self.get_param("confidence", 0.001)
        self.nms_thresh = self.get_param("nms_thresh", 0.60)
        self.imgsz = self.get_param("imgsz", 640)
        self.dataset_name = self.get_param("dataset_name", "coco_detection")
        self.pt_model = self.get_param("pt_model", "yolov11x.pt")
        self.objs_in_meter = self.get_param("objs_in_meter", {"person": [-1, 1.8], "keyboard": [0.43, -1]})  # {category_name: [w, h], ...}
        self.use_shm = self.get_param("use_shm", -1)
        self.params_help()

        self.b_use_shm = False
        if self.use_shm == 1 or (self.use_shm == -1 and platform.system() == 'Linux'):
            self.b_use_shm = True

        if self.pt_model.startswith("sms::"):
            self.local_pt_model = download_model(self.__class__.__name__, self.pt_model)
            assert self.local_pt_model is not None
        else:
            self.local_pt_model = self.pt_model

        input_url = '/' + job_name + '/sensor/image_raw'
        if len(self.specified_input_topic) > 0:
            input_url = self.specified_input_topic

        output_url = '/' + job_name + '/detector/results'
        if len(self.specified_output_topic) > 0:
            output_url = self.specified_output_topic
        
        calib_url = '/' + job_name + '/sensor/calibration_info'

        self.job_queue = Queue()
        self.queue_pool.append(self.job_queue)

        self.calib_width, self.calib_height = -1, -1
        self.camera_matrix = [712.12, 0,645.23, 0, 705.87, 327.34, 0, 0, 1]
        self.camera_matrix = np.array(self.camera_matrix).reshape(3, 3)
        self.distortion = [0.0, 0.0, 0.0, 0.0, 0.0]
        self.distortion = np.array(self.distortion)

        self._image_reader = Subscriber(
            input_url, 'std_msgs::Null', self.image_callback,
            ip=ip, port=port
        )
        self._calibration_reader = Subscriber(
            calib_url, 'sensor_msgs::CameraCalibration', self.calibration_callback,
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._result_writer = Publisher(
            output_url, 'spirecv_msgs::2DTargets',
            ip=self.remote_ip, port=self.remote_port, qos=QoS.Reliability
        )
        self._show_writer = Publisher(
            '/' + job_name + '/detector/image_results', 'memory_msgs::RawImage' if self.b_use_shm else 'sensor_msgs::CompressedImage',
            ip=ip, port=port
        )
        if self.launch_next_emit:
            self._next_writer = Publisher(
                '/' + job_name + '/launch_next', 'std_msgs::Boolean',
                ip=ip, port=port, qos=QoS.Reliability
            )

        self._detector = YOLO(self.local_pt_model)  # load an official model
        self.infer_n_ims = 0
        self.infer_time3 = 0.0
        self.infer_time2 = 0.0
        self.infer_time1 = 0.0

        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_reader.kill()
        self._calibration_reader.kill()
        self._result_writer.kill()
        self._show_writer.kill()
        self._next_writer.kill()

    def image_callback(self, msg):
        if self.realtime_det:
            if not self.job_queue.empty():
                self.job_queue.queue.clear()
        img = sms2cvimg(msg)
        self.job_queue.put({'msg': msg, 'img': img})

    def calibration_callback(self, msg):
        self.calib_width = msg['width']
        self.calib_height = msg['height']

        self.camera_matrix = np.array(msg['K']).reshape(3, 3)
        self.distortion = np.array(msg['D'])

    def run(self):
        while self.is_running():
            msg_dict = self.job_queue.get(block=True)
            if msg_dict is None:
                break

            msg, img = msg_dict['msg'], msg_dict['img']

            t1 = time.time()
            file_name = msg['file_name'] if 'file_name' in msg else ''

            self.infer_n_ims += 1
            self.infer_time1 += time.time() - t1
            # print('avg-infer time1: {}'.format(self.infer_time1 / self.infer_n_ims))

            t1 = time.time()

            if "rois" in msg and len(msg["rois"]) > 0:
                roi = msg["rois"][0]
                img_infer = img[roi[1]: roi[3], roi[0]: roi[2], :]
            else:
                roi = None
                img_infer = img

            # DO Object Detection
            results = self._detector(img_infer, conf=self.confidence, iou=self.nms_thresh, imgsz=self.imgsz, verbose=False)
            res_msg = trans_det_results(results[0], self.camera_matrix, [self.calib_width, self.calib_height], [img.shape[1], img.shape[0]], self.objs_in_meter, roi)
            res_msg['file_name'] = file_name
            res_msg['dataset'] = self.dataset_name
            if 'client_id' in msg:
                res_msg['client_id'] = msg['client_id']
            if 'file_name' in msg:
                res_msg['file_name'] = msg['file_name']
            if 'img_id' in msg:
                res_msg['img_id'] = msg['img_id']
            if 'img_total' in msg:
                res_msg['img_total'] = msg['img_total']
            res_msg['time_used'] = time.time() - t1

            self.infer_time2 += time.time() - t1
            # print('avg-infer time2: {}'.format(self.infer_time2 / self.infer_n_ims))

            t1 = time.time()

            self._result_writer.publish(res_msg)

            if 'img_total' in msg and self.launch_next_emit:
                next_msg = def_msg('std_msgs::Boolean')
                next_msg['data'] = True
                self._next_writer.publish(next_msg)

            if self.b_use_shm:
                msg = self._show_writer.cvimg2sms_mem(img)
            msg['spirecv_msgs::2DTargets'] = res_msg
            self._show_writer.publish(msg)

            self.infer_time3 += time.time() - t1
            # print('avg-infer time3: {}, totally {} imgs'.format(self.infer_time3 / self.infer_n_ims, self.infer_n_ims))
            # END

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name', '-j',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    # args = parser.parse_args()
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = YOLOv11DetNode_Cuda(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.join()
